# coding: utf-8
from __future__ import print_function
import argparse

import torch
from torch.autograd import Variable

import logging
import pickle
import codecs
import numpy as np
import nltk
import traceback
from client import Client
import json
import os
import logging

def fill_in_blank(question_json, answer_index):
    answers = question_json['choices'][answer_index]
    question = question_json['question']
    for answer in answers:
        question = question.replace('<BLANK>', answer)
    return question

class Args(object):
    question = ''
    model_dir = ''
    cuda = False


def process_args(args):
    args.dict_path = os.path.join(args.model_dir, 'dict.pkl')
    args.model_path = os.path.join(args.model_dir, 'model.pt')
    return args


def acc_score(args):
    c = Client(args)
    all_count = 0
    correct_count = 0
    with codecs.open(args.question, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            question_json = json.loads(line)
            pred_index = 0
            pred_score = -100000000
            for answer_index in range(len(question_json['choices'])):
                filled_sent = fill_in_blank(question_json, answer_index)
                _, _, _, score = c.clear_and_score(filled_sent)
                if score > pred_score:
                    pred_score = score
                    pred_index = answer_index
            all_count += 1
            if pred_index == question_json['answer']:
                correct_count += 1
            if (i + 1) % 500 == 0:
                logging.info('num: {:^10} acc: {:^10}'.format(all_count, float(correct_count) / all_count))
    acc = float(correct_count) / all_count
    logging.info('acc: {}'.format(acc))
    return acc

if __name__ == '__main__':
    args = Args()
    names = ['primary', 'middle', 'high']
    model_dirs = [
        'models/primary/emsize1000nhid1000dropout0.50tiedTrue',
        'models/middle/emsize1000nhid1000dropout0.50tiedTrue',
        'models/high/emsize1500nhid1500dropout0.50tiedTrue',
    ]
    questions = ['../data/primary/test.json', '../data/middle/test.json', '../data/high/test.json']
    for i in range(len(names)):
        for j in range(len(names)):
            args.model_dir = model_dirs[i]
            args.question = questions[j]
            process_args(args)
            acc = acc_score(args)
            print("model {} -> question {}  acc: {}".format(names[i], names[j], acc))







